import argparse
import sys
# from vision.core.imagedb import ImageDB
from dataloader.imagedb import ImageDB
from dataloader.heils_dataset import HeilsDataset, TrainAugmentation, TestTransform
from vision.train_net.train import train_pnet
import vision.config as config
import os
import random
from torch.utils.data import DataLoader
import cv2



def train_net(annotation_file, model_store_path,resume = None,
                end_epoch=16, frequent=200, lr=0.002, batch_size=512, num_workers = 8, use_cuda=True):

    cv2.setNumThreads(0)
    imagedb = ImageDB(annotation_file)
    train_imdb, test_imdb = imagedb.load_imdb()
    train_imdb = imagedb.append_flipped_images(train_imdb)

    train_transform = TrainAugmentation(size=12)
    test_transform = TestTransform(size=12)
    print('Train images number', len(train_imdb))
    print('Test images number', len(test_imdb))
    train_dataset = HeilsDataset(train_imdb, train_transform, None, is_test = False)
    test_dataset = HeilsDataset(test_imdb, None, test_transform, is_test=True)
    train_loader = DataLoader(train_dataset, batch_size,
                              num_workers=num_workers,
                              shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size,
                              num_workers=num_workers,
                              shuffle=False)
    train_pnet(model_store_path=model_store_path, resume = resume, end_epoch=end_epoch, train_loader = train_loader, test_loader=test_loader, frequent=frequent, base_lr=lr, use_cuda=use_cuda)

def parse_args():
    parser = argparse.ArgumentParser(description='Train PNet',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)


    parser.add_argument('--anno_file', dest='annotation_file',
                        default=os.path.join(config.ANNO_STORE_DIR,config.PNET_TRAIN_IMGLIST_FILENAME), help='training data annotation file', type=str)
    parser.add_argument('--model_path', dest='model_store_path', help='training model store directory',
                        default=config.PNET_MODEL_STORE_DIR, type=str)
    parser.add_argument('--end_epoch', dest='end_epoch', help='end epoch of training',
                        default=config.END_EPOCH, type=int)
    parser.add_argument('--frequent', dest='frequent', help='frequency of logging',
                        default=config.FREQUENCE, type=int)
    parser.add_argument('--lr', dest='lr', help='learning rate',
                        default=config.TRAIN_LR, type=float)
    parser.add_argument('--batch_size', dest='batch_size', help='train batch size',
                        default=config.TRAIN_BATCH_SIZE, type=int)
    parser.add_argument('--num_workers', dest='num_workers', help='torch dataloader num workers',
                        default=config.NUM_WORKERS, type=int)
    parser.add_argument('--gpu', dest='use_cuda', help='train with gpu',
                        default=config.USE_CUDA, type=bool)
    parser.add_argument('--prefix_path', dest='', help='training data annotation images prefix root path', type=str)
    parser.add_argument('--resume', dest='resume', default='', help='training data annotation images prefix root path', type=str)


    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    print('train Pnet argument:')
    print(args)

    train_net(annotation_file=args.annotation_file, model_store_path=args.model_store_path,resume = args.resume,
                end_epoch=args.end_epoch, frequent=args.frequent, lr=args.lr, batch_size=args.batch_size, num_workers = args.num_workers, use_cuda=args.use_cuda)
